use anyhow::{Result, anyhow};

/// 多测站单通道矩阵
pub struct StationDataMatrix {
    station_count: usize,
    max_sample_count: usize,
    data: Vec<f32>,
    expected_sample_counts: Vec<usize>,
}

impl StationDataMatrix {
    pub fn new(station_count: usize, max_sample_count: usize) -> Result<Self> {
        if station_count == 0 {
            return Err(anyhow!("测站数量必须大于 0"));
        }
        if max_sample_count == 0 {
            return Err(anyhow!("最大样本数必须大于 0"));
        }

        Ok(Self {
            station_count,
            max_sample_count,
            data: vec![0.0; station_count * max_sample_count],
            expected_sample_counts: vec![0; station_count],
        })
    }

    #[inline]
    fn offset(&self, station_index: usize, column_index: usize) -> Result<usize> {
        if station_index >= self.station_count {
            return Err(anyhow!(
                "测站索引越界: {} / {}",
                station_index,
                self.station_count
            ));
        }
        if column_index >= self.max_sample_count {
            return Err(anyhow!(
                "样本索引越界: {} / {}",
                column_index,
                self.max_sample_count
            ));
        }
        Ok(station_index * self.max_sample_count + column_index)
    }

    pub fn write_sample(
        &mut self,
        station_index: usize,
        column_index: usize,
        value: f32,
    ) -> Result<()> {
        let idx = self.offset(station_index, column_index)?;
        self.data[idx] = value;
        if self.expected_sample_counts[station_index] < column_index + 1 {
            self.expected_sample_counts[station_index] = column_index + 1;
        }
        Ok(())
    }

    pub fn station_count(&self) -> usize {
        self.station_count
    }

    pub fn max_sample_count(&self) -> usize {
        self.max_sample_count
    }

    pub fn data(&self) -> &[f32] {
        &self.data
    }

    pub fn data_mut(&mut self) -> &mut [f32] {
        &mut self.data
    }

    pub fn expected_sample_counts(&self) -> &[usize] {
        &self.expected_sample_counts
    }
}
